class VYM():

    def __init__(self,file1="vasprun.xml",file2="layer"):
        import numpy as np
        self.D2,self.latt,self.pos=self.__obtain_FC__(file1)
        self.layer,self.thick=self.__obtain_layer__(file2)
        if self.thick <=0:
            top_atom,bottom_atom=self.__findthick__()
            self.thick = (self.pos[self.layer[0][top_atom],2]-self.pos[self.layer[-1][bottom_atom],2])
        s=[self.latt[0,1]*self.latt[1,2]-self.latt[0,2]*self.latt[1,1],self.latt[0,2]*self.latt[1,0]-self.latt[1,2]*self.latt[0,0],self.latt[0,0]*self.latt[1,1]-self.latt[0,1]*self.latt[1,0]]
        self.area=np.linalg.norm(s)

    def __obtain_FC__(self,file):
        import xml.etree.ElementTree as ET
        import numpy as np
        tree = ET.parse(file)
        root = tree.getroot()
        mark = True
        n=0
        mass=[]
        latt=np.mat(np.zeros((3,3)))
        while mark:
            sub1 = root[n]
            n += 1
            for sub2 in sub1:
                if sub2.tag == 'atoms':
                    natom = int(sub2.text)
                if sub2.tag == 'types':
                    ntypes = int(sub2.text)
                    types = [[0,0] for i in range(ntypes)]
                if "atomtypes" in  sub2.attrib.values():
                    for sub3 in sub2:
                        if sub3.tag == 'set':
                            for i in range(ntypes):
                                types[i][0] =int(sub3[i][0].text)
                                types[i][1] =float(sub3[i][2].text)
                if sub2.tag == "crystal":
                    for sub3 in sub2:
                        if "basis" in  sub3.attrib.values():
                            latt[0,:]=np.mat(sub3[0].text)
                            latt[1,:]=np.mat(sub3[1].text)
                            latt[2,:]=np.mat(sub3[2].text)
                if "positions" in  sub2.attrib.values():
                    atom_pos=np.mat(np.zeros((natom,3)))
                    for i in range(natom):
                        atom_pos[i,:]=np.mat(sub2[i].text)
                if sub2.tag == 'dynmat':
                    dynMatrix=np.mat(np.zeros((3*natom,3*natom)))
                    for sub3 in sub2:
                        if "hessian" in sub3.attrib.values():
                            for atom in range(3*natom):
                                  dynMatrix[atom,:]=np.mat(sub3[atom].text)
                    mark = False
        for i in range(ntypes):
            for j in range(types[i][0]):
                mass.append(types[i][1])
                mass.append(types[i][1])
                mass.append(types[i][1])
        mass_matrix = np.sqrt(np.multiply(np.mat(mass),np.transpose(np.mat(mass))))
        D2 = np.multiply(dynMatrix,mass_matrix)
        return D2,latt,atom_pos*latt

    def __obtain_layer__(self,file2):
        fid=open(file2,"rt")
        lines=fid.readlines()
        nlayer = len(lines)-1
        layer=[]
        for i in range(nlayer):
            layer.append([])
            for atom in lines[i].split(','):
                layer[i].append(int(atom)-1)
        thick=float(lines[-1])
        return layer,thick

    def __findthick__(self):
        top_pos=-1e3
        bottom_pos=1e3
        for i in range(len(self.layer[0])):
            if self.pos[self.layer[0][i],2]>top_pos:
                top_pos=self.pos[self.layer[0][i],2]
                top_atom=i
        for i in range(len(self.layer[-1])):
            if self.pos[self.layer[-1][i],2]<bottom_pos:
                bottom_pos=self.pos[self.layer[-1][i],2]
                bottom_atom=i
        return  top_atom,bottom_atom

    def __cal_k__(self,layer1,layer2):
        k=0
        for i in range(len(layer1)):
            for j in range(len(layer2)):
                k+=self.D2[3*layer1[i]+2,3*layer2[j]+2]
        return k

    def main(self):
         k_list=[0 for i in range(int(len(self.layer)/2))]
         k_tot_inv=0
         for i in range(int(len(self.layer)/2)):
             k_list[i]=self. __cal_k__(self.layer[2*i],self.layer[2*i+1])*1.602176634*100/self.area
             k_tot_inv+=1/k_list[i]
         k_tot=1/ k_tot_inv
         YM=k_tot*self.thick
         print("Effective Thickness:\t%5.2f Ang"%self.thick)
         print("Young's Modulus:\t%5.2f Gpa"%YM)

if __name__=="__main__":
    a=VYM()
    a.main()